package com.lxz.shardingjdbcdemo.config;

import com.alibaba.druid.pool.DruidDataSource;
import org.apache.shardingsphere.driver.api.ShardingSphereDataSourceFactory;
import org.apache.shardingsphere.infra.config.algorithm.AlgorithmConfiguration;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.strategy.sharding.StandardShardingStrategyConfiguration;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.*;

/**
 * ShardingSphere分库分表配置类
 * 该类用于配置分库分表规则、数据源以及分片算法
 */
@Configuration
public class ShardingSphereConfig {

    /**
     * 创建并配置ShardingSphere数据源
     * 
     * @return 配置好的ShardingSphere数据源
     * @throws SQLException SQL异常
     */
    @Bean
    public DataSource dataSource(BroadcastTableConfig broadcastTableConfig) throws SQLException {
        // 创建数据源映射
        Map<String, DataSource> dataSourceMap = new HashMap<>();

        // 配置订单数据库 sharding_db_0
        DruidDataSource orderDataSource0 = new DruidDataSource();
        orderDataSource0.setDriverClassName("com.mysql.cj.jdbc.Driver");
        orderDataSource0.setUrl("jdbc:mysql://localhost:3306/sharding_db_0?useUnicode=true&characterEncoding=utf-8&serverTimezone=UTC");
        orderDataSource0.setUsername("root");
        orderDataSource0.setPassword("123456");
        dataSourceMap.put("ds0", orderDataSource0);

        // 配置订单数据库 sharding_db_1
        DruidDataSource orderDataSource1 = new DruidDataSource();
        orderDataSource1.setDriverClassName("com.mysql.cj.jdbc.Driver");
        orderDataSource1.setUrl("jdbc:mysql://localhost:3306/sharding_db_1?useUnicode=true&characterEncoding=utf-8&serverTimezone=UTC");
        orderDataSource1.setUsername("root");
        orderDataSource1.setPassword("123456");
        dataSourceMap.put("ds1", orderDataSource1);

        // 配置用户数据库 user_db
        DruidDataSource userDataSource = new DruidDataSource();
        userDataSource.setDriverClassName("com.mysql.cj.jdbc.Driver");
        userDataSource.setUrl("jdbc:mysql://localhost:3306/user_db?useUnicode=true&characterEncoding=utf-8&serverTimezone=UTC");
        userDataSource.setUsername("root");
        userDataSource.setPassword("123456");
        dataSourceMap.put("user_db", userDataSource);

        // 配置分片规则
        ShardingRuleConfiguration shardingRuleConfig = new ShardingRuleConfiguration();
        shardingRuleConfig.getTables().add(getOrderTableRuleConfiguration());
        shardingRuleConfig.getTables().add(getUserTableRuleConfiguration());
        // 移除t_user_type的分片规则配置，将其作为广播表处理
        shardingRuleConfig.getShardingAlgorithms().put("database-inline", createDBInlineShardingAlgorithmConfiguration());
        shardingRuleConfig.getShardingAlgorithms().put("table-inline", createOrderInlineShardingAlgorithmConfiguration());
        shardingRuleConfig.getKeyGenerators().put("snowflake", createKeyGeneratorConfiguration());

        // 配置属性
        Properties props = new Properties();
        props.setProperty("sql-show", "true");

        // 创建并返回ShardingSphere数据源，包含分片规则和广播表规则
        return ShardingSphereDataSourceFactory.createDataSource(
            dataSourceMap, 
            Arrays.asList(shardingRuleConfig, broadcastTableConfig.broadcastRuleConfiguration()), 
            props
        );
    }

    /**
     * 获取订单表的分片规则配置
     * 
     * @return 订单表分片规则配置
     */
    private ShardingTableRuleConfiguration getOrderTableRuleConfiguration() {
        // 配置逻辑表名和实际数据节点
        ShardingTableRuleConfiguration result = new ShardingTableRuleConfiguration("t_order", "ds${0..1}.t_order_${0..1}");
        
        // userId作为分库分片键
        result.setDatabaseShardingStrategy(new StandardShardingStrategyConfiguration("user_id", "database-inline"));
        
        // orderId作为分表分片键
        result.setTableShardingStrategy(new StandardShardingStrategyConfiguration("order_id", "table-inline"));
        
        // 设置主键生成器
        result.setKeyGenerateStrategy(new org.apache.shardingsphere.sharding.api.config.strategy.keygen.KeyGenerateStrategyConfiguration("order_id", "snowflake"));
        
        return result;
    }
    
    /**
     * 获取用户表的分片规则配置
     * 
     * @return 用户表分片规则配置
     */
    private ShardingTableRuleConfiguration getUserTableRuleConfiguration() {
        // 配置用户表规则，指定数据源为user_db
        ShardingTableRuleConfiguration result = new ShardingTableRuleConfiguration("t_user", "user_db.t_user");
        return result;
    }

    /**
     * 创建分库分片算法配置
     * 
     * @return 分库分片算法配置
     */
    private AlgorithmConfiguration createDBInlineShardingAlgorithmConfiguration() {
        Properties props = new Properties();
        // 使用user_id对2取模来决定数据存储在哪个数据库
        props.setProperty("algorithm-expression", "ds${user_id % 2}");
        return new AlgorithmConfiguration("INLINE", props);
    }

    /**
     * 创建分表分片算法配置
     * 
     * @return 分表分片算法配置
     */
    private AlgorithmConfiguration createOrderInlineShardingAlgorithmConfiguration() {
        Properties props = new Properties();
        // 使用order_id对2取模来决定数据存储在哪个表
        props.setProperty("algorithm-expression", "t_order_${order_id % 2}");
        return new AlgorithmConfiguration("INLINE", props);
    }
    
    /**
     * 创建主键生成器配置
     * 
     * @return 主键生成器配置
     */
    private AlgorithmConfiguration createKeyGeneratorConfiguration() {
        Properties props = new Properties();
        // 设置雪花算法的工作节点ID
        props.setProperty("worker-id", "123");
        return new AlgorithmConfiguration("SNOWFLAKE", props);
    }
}